"""
YOLOv8 提供多种预训练模型（从轻量型到高精度），可从官方仓库下载，或直接通过代码自动下载。常见模型：
yolov8n.pt（nano，轻量快）
yolov8s.pt（small，平衡）
yolov8m.pt（medium，较高精度）
yolov8l.pt/yolov8x.pt（large/xlarge，高精度但慢）
"""

from ultralytics import YOLO
import cv2
import os
from PIL import Image

def dataset_generate_auto():
    """
    自动标注
    
    """
    # 加载预训练模型
    model = YOLO("yolov8s.pt")  # 用COCO预训练模型

    # 无标注图像的文件夹路径
    image_dir = "path/to/unlabeled_images"
    # 生成的标注文件保存路径（需与图像同名，放在labels文件夹）
    label_dir = "path/to/labels"
    os.makedirs(label_dir, exist_ok=True)

    # 对所有图像推理并生成标注
    for img_name in os.listdir(image_dir):
        if img_name.endswith((".jpg", ".png")):
            img_path = os.path.join(image_dir, img_name)
            results = model(img_path)  # 推理
            
            # 提取标注信息并保存为YOLO格式的txt文件
            label_path = os.path.join(label_dir, img_name.replace(".jpg", ".txt").replace(".png", ".txt"))
            with open(label_path, "w") as f:
                for box in results[0].boxes:
                    cls = int(box.cls)  # 类别ID（注意：预训练模型的类别可能与你的任务不同，需映射）
                    conf = box.conf.item()  # 置信度
                    if conf > 0.5:  # 只保留置信度>0.5的预测
                        # 转换为归一化坐标（x_center, y_center, width, height）
                        x1, y1, x2, y2 = box.xyxy[0]  # 左上角和右下角坐标（像素值）
                        img_h, img_w = results[0].orig_shape  # 图像原始高宽
                        x_center = (x1 + x2) / 2 / img_w
                        y_center = (y1 + y2) / 2 / img_h
                        width = (x2 - x1) / img_w
                        height = (y2 - y1) / img_h
                        f.write(f"{cls} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

def dataset_generate_manual(image_dir,cls_id:int):
    """
    手动标注
    若坐标是像素值（未归一化），需按公式转换为 0~1 范围
    """
    # 生成的标注文件保存路径（需与图像同名，放在labels文件夹）
    label_dir = os.path.join(image_dir, "../labels")
    os.makedirs(label_dir, exist_ok=True)

    # 对所有图像推理并生成标注
    for img_name in os.listdir(image_dir):
        if img_name.endswith((".jpg", ".png")):
            img_path = os.path.join(image_dir, img_name)
            img = Image.open(img_path)
            w,h = img.size
            # print(w,h)
            #提取标注信息并保存为YOLO格式的txt文件
            label_path = os.path.join(label_dir, img_name.replace(".jpg", ".txt").replace(".png", ".txt"))
            with open(label_path, "w") as f:
                # 转换为归一化坐标（x_center, y_center, width, height）
                x1,y1,x2,y2 = 0,0,w,h        #这里将整张照片多为目标区域，所以坐标为0,0,w,h
                x_center = (x1 + x2) / 2 / w
                y_center = (y1 + y2) / 2 / h
                width = (x2 - x1) / w
                height = (y2 - y1) / h
                f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
    print(f"标注完成:{image_dir}")




def train_yolo_model():
    # 加载预训练模型
    model = YOLO(r"D:\yolo11_base\yolo11s.pt")  # 可替换为其他预训练模型

    # 微调训练
    results = model.train(
        data="./6项目实战\大模型微调\custom_data.yaml",  # 数据集配置文件路径
        epochs=100,                # 训练轮次（根据数据集大小调整，通常30~100）
        batch=16,                 # 批次大小（根据GPU显存调整，如8/16/32）
        imgsz=640,                # 输入图像尺寸（YOLO默认640x640）
        device=0,                 # 训练设备（0为GPU，-1为CPU）
        pretrained=True,          # 启用预训练权重（默认True）
        freeze=10,                # 冻结前10层（先训练头部，再解冻全局）
        lr0=0.001                 # 初始学习率（微调通常较小，默认0.01）
    )

"""
关键参数说明
freeze：冻结主干网络（backbone）的前 N 层，仅训练头部（head）。目的是先让模型适应新数据集的类别，避免破坏预训练的特征提取能力。通常先冻结训练 10~20 轮，再解冻（freeze=0）继续训练。
epochs：若数据集小（如几千张），50~100 轮足够；若数据量大（几万张），可增加到 200 轮以上。
batch：受 GPU 显存限制，显存不足时减小（如 8），足够时增大（如 32）以提高训练效率。
data_augmentation：YOLOv8 默认启用数据增强（如随机裁剪、翻转、缩放等），可通过augment=True加强，避免过拟合。
"""

def detect_signature(image_path):
    # 加载训练好的模型
    model = YOLO(r"./runs\detect\train2\weights\best.pt")  # 替换为你的模型路径

    # 对单张图像进行推理
    results = model(image_path)    # 执行推理
    result = results[0]
    print("预测框坐标：", result.boxes.xyxy)  # 左上角x、y，右下角x、y（像素值）
    print("置信度：", result.boxes.conf)     # 每个预测框的置信度
    print("类别ID：", result.boxes.cls)      # 每个预测框的类别ID

    # 处理并可视化结果
    for result in results:
        # 保存带预测框的图像（默认保存在 runs/detect/predict 目录）
        result.save(filename="predicted_image.jpg")
        # 也可以用OpenCV显示图像
        img = result.plot()  # 获取带预测框的图像（BGR格式）
        cv2.imshow("Prediction", img)
        cv2.waitKey(0)  # 按任意键关闭窗口
        cv2.destroyAllWindows()

if __name__ == "__main__":
    # train_yolo_model()
    detect_signature(r"D:\Signature/5.jpeg")
    # dataset_generate_manual(r"D:\Garbage_Classification_YOLO\classify\val\trash",cls_id=9)

